from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.cm as cm

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def get_intrinsic_pc():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        diff = jnp.abs(rnn_state - rnn_state_old)
        rnn_state_old = rnn_state.copy()
        diff_norm = jnp.linalg.norm(diff, axis=1)
        norm_std = 0.4 * jnp.abs(diff_norm - diff_norm_old) + 0.6 * norm_std
        diff_norm_old = diff_norm.copy()

        rnn_state_trajectory.append(np.array(rnn_state).copy())
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 对 rnn_state_np 进行 PCA
    pca = PCA()
    # pca.fit(rnn_state_np)
    pca.fit(rnn_state_trajectory_np)

    # 打印 variance ratio
    print(pca.explained_variance_ratio_)

    rnn_state_np_pca = pca.transform(rnn_state_np)

    # 创建KMeans对象，指定聚类数为4
    kmeans = KMeans(n_clusters=4)
    # 对rnn_state_np_pca进行聚类
    kmeans.fit(rnn_state_np_pca)
    # 获取聚类中心的坐标
    cluster_centers = kmeans.cluster_centers_

    return pca, cluster_centers


def main():
    
    seq_len = 15
    
    n_samples = 1000

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--seq_len", type=int, default=15)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    seq_len = args.seq_len

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    
    arrow_length = 1
    arrow_list = np.array([[arrow_length, 0], [-arrow_length, 0], [0, arrow_length], [0, -arrow_length], [0, 0]])

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    # 生成一个二进制串集合，其中每个元素都是一个长度为8的二进制串，要求这些二进制串从 000000000 遍历到 111111111
    binary_set = set()
    for i in range(256):
        binary_string = format(i, '08b')
        binary_set.add(binary_string)
    # 将 binary_set 中的 11111111 元素删除
    binary_set.remove("11111111")
    # 将 binary_set 中所有格式为 "x1x1x1x1" 的元素删除
    binary_set_bk = binary_set.copy()
    for binary in binary_set_bk:
        if binary[1] == '1' and binary[3] == '1' and binary[4] == '1' and binary[6] == '1':
            binary_set.remove(binary)
    print("binary_set: ", binary_set)
    
    # 对 binary_set 进行排序，使得其中的元素从 00000000 遍历到 11111110
    binary_list = sorted(binary_set)
    # 将 binary_list 中的每个元素进行这样的操作：在中间插入一个0，例如 "00000000" 转换为 "000000000"；然后在结尾处插入一个0，例如 "000000000" 转换为 "0000000000"
    binary_list = [i[:4] + "0" + i[4:] + "0" for i in binary_list]
    # 将 binary_list 中的元素转换为整数数组，例如 "00000000" 转换为 [0, 0, 0, 0, 0, 0, 0, 0]
    binary_list = [list(map(int, list(i))) for i in binary_list]
    binary_list = np.array(binary_list)
    print("shape of binary_list: ", binary_list.shape)

    binary_list, arrow_list = np.array(binary_list), np.array(arrow_list)

    def make_obs_img(obs_int):
        # 将形状为 (10,) 的 obs 裁减掉最后一位，变成形状为 (9,) 的 obs
        obs = obs_int[:-1]
        # 将 obs 转换为形状为 (3, 3) 的 numpy 数组
        obs = obs.reshape((3, 3))
        
        # 将 obs 中的 1 替换为 255
        obs = (1-obs) * 255
        # 将 obs 转换为形状为 (3, 3, 1) 的 numpy 数组
        obs = obs.reshape((3, 3, 1))
        # 将 obs 转换为形状为 (3, 3, 3) 的 numpy 数组
        obs = np.concatenate((obs, obs, obs), axis=2)
        # 将 obs 转换为 opencv 的8比特图像格式
        obs = obs.astype(np.uint8)
        # 将 obs 转换为形状为 (60, 60, 3) 的 numpy 数组
        obs = cv2.resize(obs, (60, 60), interpolation=cv2.INTER_NEAREST)
        # 在 obs 上绘制 3x3 的灰色网格
        for i in range(1, 3):
            cv2.line(obs, (0, i*20), (60, i*20), (100, 100, 100), 1)
            cv2.line(obs, (i*20, 0), (i*20, 60), (100, 100, 100), 1)
        # 在 obs 的边缘绘制灰色边框
        cv2.rectangle(obs, (0, 0), (59, 59), (100, 100, 100), 1)
        return obs

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ load task trajectories
    """
    slice_fn = "_" + nn_type + "_" + str(seq_len)
    
    # 载入随机的切片数据
    trajectories_diverse_set = np.load("./logs/er9_data/trial2/tracelet_slices"+slice_fn+".npz")
    obs_record_diverse_set = np.load("./logs/er9_data/trial2/obs_slices"+slice_fn+".npz")
    action_slices_diverse_set = np.load("./logs/er9_data/trial2/action_slices"+slice_fn+".npz")

    tracelet_slices1 = trajectories_diverse_set["mat1"]
    tracelet_slices2 = trajectories_diverse_set["mat2"]
    tracelet_slices3 = trajectories_diverse_set["mat3"]
    tracelet_slices4 = trajectories_diverse_set["mat4"]
    tracelet_slices5 = trajectories_diverse_set["mat5"]

    obs_slices1 = obs_record_diverse_set["mat1"]
    obs_slices2 = obs_record_diverse_set["mat2"]
    obs_slices3 = obs_record_diverse_set["mat3"]
    obs_slices4 = obs_record_diverse_set["mat4"]
    obs_slices5 = obs_record_diverse_set["mat5"]

    action_slices1 = action_slices_diverse_set["mat1"]
    action_slices2 = action_slices_diverse_set["mat2"]
    action_slices3 = action_slices_diverse_set["mat3"]
    action_slices4 = action_slices_diverse_set["mat4"]
    action_slices5 = action_slices_diverse_set["mat5"]

    print("shape of tracelet_slices1: ", tracelet_slices1.shape)
    print("shape of tracelet_slices2: ", tracelet_slices2.shape)
    print("shape of tracelet_slices3: ", tracelet_slices3.shape)
    print("shape of tracelet_slices4: ", tracelet_slices4.shape)
    print("shape of tracelet_slices5: ", tracelet_slices5.shape)

    print("shape of obs_slices1: ", obs_slices1.shape)
    print("shape of obs_slices2: ", obs_slices2.shape)
    print("shape of obs_slices3: ", obs_slices3.shape)
    print("shape of obs_slices4: ", obs_slices4.shape)
    print("shape of obs_slices5: ", obs_slices5.shape)

    print("shape of action_slices1: ", action_slices1.shape)
    print("shape of action_slices2: ", action_slices2.shape)
    print("shape of action_slices3: ", action_slices3.shape)
    print("shape of action_slices4: ", action_slices4.shape)
    print("shape of action_slices5: ", action_slices5.shape)


    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    """ action list
    """
    action_list = jnp.array([[0, 1], [0, -1], [1, 0], [-1, 0], [0, 0]])


    """ 遍历所有轨迹片段, 收集每一条obs序列的最小 action 误差极限环
        要求最小 action 误差小于等于2
        收集满7条合规轨迹, 或者遍历完所有轨迹片段后结束
    """

    collect_capacity = 15
    collect_err_th = 3

    rnn_state_trajectory_set1 = []
    rnn_state_trajectory_set2 = []
    rnn_state_trajectory_set3 = []
    rnn_state_trajectory_set4 = []
    rnn_state_trajectory_set5 = []

    for s in range(tracelet_slices1.shape[0]):
        progress_bar(len(rnn_state_trajectory_set1), collect_capacity)

        obs_record = obs_slices1[s]
        action_slice = action_slices1[s]

        obs_seq_binary = []
        for obs in obs_record:
            # 查找 obs 在 binary_list 中的序号
            index_ = np.argwhere(np.all(obs == binary_list, axis=1))
            obs_seq_binary.append(index_[0][0])
        
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []

        random_integers = obs_seq_binary

        random_integers_i = 0
        def update_random_integers_i():
            nonlocal random_integers_i
            nonlocal seq_len
            random_integers_i += 1
            random_integers_i %= seq_len
        
        tracelets = []
        tracelet = []
        actionlets = []
        actionlet = []

        """ run forced dynamics
        """
        for t in range(rpl_config.life_duration):

            tracelet.append(np.array(rnn_state).copy())
            if random_integers_i == seq_len-1:
                tracelets.append(tracelet.copy())
                tracelet.clear()

            random_obs = binary_list[random_integers[random_integers_i]]
            obs_zero = jnp.array([random_obs for i in range(n_samples)])

            """ model forward 
            """
            rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            actions = get_action_vmap(y1)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            update_random_integers_i()

            actionlet.append(actions.copy())
            if random_integers_i == seq_len-1:
                actionlets.append(actionlet.copy())
                actionlet.clear()

        actionlets = actionlets[1:-2]
        actionlets = np.array(actionlets)
        # 将 actionlets 的第2维和第3维交换
        actionlets = np.swapaxes(actionlets, 1, 2)
        print("shape of actionlets: ", actionlets.shape)

        actionlets_limit_cycle = actionlets[-2]
        print("shape of actionlets_limit_cycles: ", actionlets_limit_cycle.shape)

        # 将 rnn_state_trajectory 的第1、2维交换
        rnn_state_trajectory = np.array(rnn_state_trajectory)
        rnn_state_trajectory = np.swapaxes(rnn_state_trajectory, 0, 1)

        """ 逐个检测 action 序列和原始数据的相似度
        """
        diff = []

        for i in range(actionlets_limit_cycle.shape[0]):
            actionlet = actionlets_limit_cycle[i]

            # 计算 actionlet 和 action_slice 中不同元素的个数
            diff_n_elements = np.sum(actionlet != action_slice)

            diff.append(diff_n_elements)
        diff = np.array(diff)
        print("min diff of set1: ", np.min(diff))

        min_diff = np.min(diff)
        minimum_index = np.argmin(diff)

        if min_diff <= collect_err_th:
            rnn_state_trajectory_set1.append(rnn_state_trajectory[minimum_index].copy())
            print("shape of rnn_state_trajectory_set1: ", np.array(rnn_state_trajectory_set1).shape)
            if len(rnn_state_trajectory_set1) >= collect_capacity:
                break

    for s in range(tracelet_slices2.shape[0]):
        progress_bar(len(rnn_state_trajectory_set2), collect_capacity)
        
        obs_record = obs_slices2[s]
        action_slice = action_slices2[s]

        obs_seq_binary = []
        for obs in obs_record:
            # 查找 obs 在 binary_list 中的序号
            index_ = np.argwhere(np.all(obs == binary_list, axis=1))
            obs_seq_binary.append(index_[0][0])
        
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []

        random_integers = obs_seq_binary

        random_integers_i = 0
        def update_random_integers_i():
            nonlocal random_integers_i
            nonlocal seq_len
            random_integers_i += 1
            random_integers_i %= seq_len
        
        tracelets = []
        tracelet = []
        actionlets = []
        actionlet = []

        """ run forced dynamics
        """
        for t in range(rpl_config.life_duration):

            tracelet.append(np.array(rnn_state).copy())
            if random_integers_i == seq_len-1:
                tracelets.append(tracelet.copy())
                tracelet.clear()

            random_obs = binary_list[random_integers[random_integers_i]]
            obs_zero = jnp.array([random_obs for i in range(n_samples)])

            """ model forward 
            """
            rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            actions = get_action_vmap(y1)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            update_random_integers_i()

            actionlet.append(actions.copy())
            if random_integers_i == seq_len-1:
                actionlets.append(actionlet.copy())
                actionlet.clear()

        actionlets = actionlets[1:-2]
        actionlets = np.array(actionlets)
        # 将 actionlets 的第2维和第3维交换
        actionlets = np.swapaxes(actionlets, 1, 2)
        print("shape of actionlets: ", actionlets.shape)

        actionlets_limit_cycle = actionlets[-2]
        print("shape of actionlets_limit_cycles: ", actionlets_limit_cycle.shape)

        # 将 rnn_state_trajectory 的第1、2维交换
        rnn_state_trajectory = np.array(rnn_state_trajectory)
        rnn_state_trajectory = np.swapaxes(rnn_state_trajectory, 0, 1)

        """ 逐个检测 action 序列和原始数据的相似度
        """
        diff = []

        for i in range(actionlets_limit_cycle.shape[0]):
            actionlet = actionlets_limit_cycle[i]

            # 计算 actionlet 和 action_slice 中不同元素的个数
            diff_n_elements = np.sum(actionlet != action_slice)

            diff.append(diff_n_elements)
        diff = np.array(diff)
        print("min diff of set1: ", np.min(diff))

        min_diff = np.min(diff)
        minimum_index = np.argmin(diff)

        if min_diff <= collect_err_th:
            rnn_state_trajectory_set2.append(rnn_state_trajectory[minimum_index].copy())
            print("shape of rnn_state_trajectory_set2: ", np.array(rnn_state_trajectory_set2).shape)
            if len(rnn_state_trajectory_set2) >= collect_capacity:
                break

    for s in range(tracelet_slices3.shape[0]):
        progress_bar(len(rnn_state_trajectory_set3), collect_capacity)
        
        obs_record = obs_slices3[s]
        action_slice = action_slices3[s]

        obs_seq_binary = []
        for obs in obs_record:
            # 查找 obs 在 binary_list 中的序号
            index_ = np.argwhere(np.all(obs == binary_list, axis=1))
            obs_seq_binary.append(index_[0][0])
        
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []

        random_integers = obs_seq_binary

        random_integers_i = 0
        def update_random_integers_i():
            nonlocal random_integers_i
            nonlocal seq_len
            random_integers_i += 1
            random_integers_i %= seq_len
        
        tracelets = []
        tracelet = []
        actionlets = []
        actionlet = []

        """ run forced dynamics
        """
        for t in range(rpl_config.life_duration):

            tracelet.append(np.array(rnn_state).copy())
            if random_integers_i == seq_len-1:
                tracelets.append(tracelet.copy())
                tracelet.clear()

            random_obs = binary_list[random_integers[random_integers_i]]
            obs_zero = jnp.array([random_obs for i in range(n_samples)])

            """ model forward 
            """
            rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            actions = get_action_vmap(y1)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            update_random_integers_i()

            actionlet.append(actions.copy())
            if random_integers_i == seq_len-1:
                actionlets.append(actionlet.copy())
                actionlet.clear()

        actionlets = actionlets[1:-2]
        actionlets = np.array(actionlets)
        # 将 actionlets 的第2维和第3维交换
        actionlets = np.swapaxes(actionlets, 1, 2)
        print("shape of actionlets: ", actionlets.shape)

        actionlets_limit_cycle = actionlets[-2]
        print("shape of actionlets_limit_cycles: ", actionlets_limit_cycle.shape)

        # 将 rnn_state_trajectory 的第1、2维交换
        rnn_state_trajectory = np.array(rnn_state_trajectory)
        rnn_state_trajectory = np.swapaxes(rnn_state_trajectory, 0, 1)

        """ 逐个检测 action 序列和原始数据的相似度
        """
        diff = []

        for i in range(actionlets_limit_cycle.shape[0]):
            actionlet = actionlets_limit_cycle[i]

            # 计算 actionlet 和 action_slice 中不同元素的个数
            diff_n_elements = np.sum(actionlet != action_slice)

            diff.append(diff_n_elements)
        diff = np.array(diff)
        print("min diff of set1: ", np.min(diff))

        min_diff = np.min(diff)
        minimum_index = np.argmin(diff)

        if min_diff <= collect_err_th:
            rnn_state_trajectory_set3.append(rnn_state_trajectory[minimum_index].copy())
            print("shape of rnn_state_trajectory_set3: ", np.array(rnn_state_trajectory_set3).shape)
            if len(rnn_state_trajectory_set3) >= collect_capacity:
                break

    for s in range(tracelet_slices4.shape[0]):
        progress_bar(len(rnn_state_trajectory_set4), collect_capacity)
        
        obs_record = obs_slices4[s]
        action_slice = action_slices4[s]

        obs_seq_binary = []
        for obs in obs_record:
            # 查找 obs 在 binary_list 中的序号
            index_ = np.argwhere(np.all(obs == binary_list, axis=1))
            obs_seq_binary.append(index_[0][0])
        
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []

        random_integers = obs_seq_binary

        random_integers_i = 0
        def update_random_integers_i():
            nonlocal random_integers_i
            nonlocal seq_len
            random_integers_i += 1
            random_integers_i %= seq_len
        
        tracelets = []
        tracelet = []
        actionlets = []
        actionlet = []

        """ run forced dynamics
        """
        for t in range(rpl_config.life_duration):

            tracelet.append(np.array(rnn_state).copy())
            if random_integers_i == seq_len-1:
                tracelets.append(tracelet.copy())
                tracelet.clear()

            random_obs = binary_list[random_integers[random_integers_i]]
            obs_zero = jnp.array([random_obs for i in range(n_samples)])

            """ model forward 
            """
            rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            actions = get_action_vmap(y1)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            update_random_integers_i()

            actionlet.append(actions.copy())
            if random_integers_i == seq_len-1:
                actionlets.append(actionlet.copy())
                actionlet.clear()

        actionlets = actionlets[1:-2]
        actionlets = np.array(actionlets)
        # 将 actionlets 的第2维和第3维交换
        actionlets = np.swapaxes(actionlets, 1, 2)
        print("shape of actionlets: ", actionlets.shape)

        actionlets_limit_cycle = actionlets[-2]
        print("shape of actionlets_limit_cycles: ", actionlets_limit_cycle.shape)

        # 将 rnn_state_trajectory 的第1、2维交换
        rnn_state_trajectory = np.array(rnn_state_trajectory)
        rnn_state_trajectory = np.swapaxes(rnn_state_trajectory, 0, 1)

        """ 逐个检测 action 序列和原始数据的相似度
        """
        diff = []

        for i in range(actionlets_limit_cycle.shape[0]):
            actionlet = actionlets_limit_cycle[i]

            # 计算 actionlet 和 action_slice 中不同元素的个数
            diff_n_elements = np.sum(actionlet != action_slice)

            diff.append(diff_n_elements)
        diff = np.array(diff)
        print("min diff of set1: ", np.min(diff))

        min_diff = np.min(diff)
        minimum_index = np.argmin(diff)

        if min_diff <= collect_err_th:
            rnn_state_trajectory_set4.append(rnn_state_trajectory[minimum_index].copy())
            print("shape of rnn_state_trajectory_set4: ", np.array(rnn_state_trajectory_set4).shape)
            if len(rnn_state_trajectory_set4) >= collect_capacity:
                break

    for s in range(tracelet_slices5.shape[0]):
        progress_bar(len(rnn_state_trajectory_set5), collect_capacity)
        
        obs_record = obs_slices5[s]
        action_slice = action_slices5[s]

        obs_seq_binary = []
        for obs in obs_record:
            # 查找 obs 在 binary_list 中的序号
            index_ = np.argwhere(np.all(obs == binary_list, axis=1))
            obs_seq_binary.append(index_[0][0])
        
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []

        random_integers = obs_seq_binary

        random_integers_i = 0
        def update_random_integers_i():
            nonlocal random_integers_i
            nonlocal seq_len
            random_integers_i += 1
            random_integers_i %= seq_len
        
        tracelets = []
        tracelet = []
        actionlets = []
        actionlet = []

        """ run forced dynamics
        """
        for t in range(rpl_config.life_duration):

            tracelet.append(np.array(rnn_state).copy())
            if random_integers_i == seq_len-1:
                tracelets.append(tracelet.copy())
                tracelet.clear()

            random_obs = binary_list[random_integers[random_integers_i]]
            obs_zero = jnp.array([random_obs for i in range(n_samples)])

            """ model forward 
            """
            rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            actions = get_action_vmap(y1)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            update_random_integers_i()

            actionlet.append(actions.copy())
            if random_integers_i == seq_len-1:
                actionlets.append(actionlet.copy())
                actionlet.clear()

        actionlets = actionlets[1:-2]
        actionlets = np.array(actionlets)
        # 将 actionlets 的第2维和第3维交换
        actionlets = np.swapaxes(actionlets, 1, 2)
        print("shape of actionlets: ", actionlets.shape)

        actionlets_limit_cycle = actionlets[-2]
        print("shape of actionlets_limit_cycles: ", actionlets_limit_cycle.shape)

        # 将 rnn_state_trajectory 的第1、2维交换
        rnn_state_trajectory = np.array(rnn_state_trajectory)
        rnn_state_trajectory = np.swapaxes(rnn_state_trajectory, 0, 1)

        """ 逐个检测 action 序列和原始数据的相似度
        """
        diff = []

        for i in range(actionlets_limit_cycle.shape[0]):
            actionlet = actionlets_limit_cycle[i]

            # 计算 actionlet 和 action_slice 中不同元素的个数
            diff_n_elements = np.sum(actionlet != action_slice)

            diff.append(diff_n_elements)
        diff = np.array(diff)
        print("min diff of set1: ", np.min(diff))

        min_diff = np.min(diff)
        minimum_index = np.argmin(diff)

        if min_diff <= collect_err_th:
            rnn_state_trajectory_set5.append(rnn_state_trajectory[minimum_index].copy())
            print("shape of rnn_state_trajectory_set5: ", np.array(rnn_state_trajectory_set5).shape)
            if len(rnn_state_trajectory_set5) >= collect_capacity:
                break

    # 将 rnn_state_trajectory_set1-5 合并成 rnn_state_trajectory_set
    rnn_state_trajectory_set = rnn_state_trajectory_set1 + rnn_state_trajectory_set2 + rnn_state_trajectory_set3 + rnn_state_trajectory_set4 + rnn_state_trajectory_set5
    rnn_state_trajectory_set = np.array(rnn_state_trajectory_set)
    rnn_state_trajectory_set0 = rnn_state_trajectory_set.copy()
    rnn_state_trajectory_ = rnn_state_trajectory_set0.reshape(-1, rnn_state_trajectory_set0.shape[-1])
    print("shape of rnn_state_trajectory_: ", rnn_state_trajectory_.shape)

    # 对 rnn_state_trajectory_ 进行 PCA
    pca = PCA()
    pca.fit(rnn_state_trajectory_)

    # 随机定义五种颜色，用于后续的 plt 绘制
    colors = ["red", "green", "blue", "orange", "purple"]

    # # 将5个集合的神经轨迹的末端极限环部分，用不同的颜色绘制到 PCA 空间里
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')

    # for i in range(5):
    #     for ii in range(collect_capacity):
    #         rnn_state_trajectory_pca = pca.transform(rnn_state_trajectory_set[i*collect_capacity + ii])
    #         ax.plot(rnn_state_trajectory_pca[-seq_len*2:-1, 0], rnn_state_trajectory_pca[-seq_len*2:-1, 1], rnn_state_trajectory_pca[-seq_len*2:-1, 2], c=colors[i], alpha=0.5)
    
    # plt.show()

    # # 用 plt 绘制 tracelet_slices1-5
    # fig, ax = plt.subplots()
    # ax.plot(tracelet_slices1[0, :, 0], tracelet_slices1[0, :, 1], c=colors[0])
    # ax.plot(tracelet_slices2[0, :, 0], tracelet_slices2[0, :, 1], c=colors[1])
    # ax.plot(tracelet_slices3[0, :, 0], tracelet_slices3[0, :, 1], c=colors[2])
    # ax.plot(tracelet_slices4[0, :, 0], tracelet_slices4[0, :, 1], c=colors[3])
    # ax.plot(tracelet_slices5[0, :, 0], tracelet_slices5[0, :, 1], c=colors[4])
    # plt.show()

    # 创建一个包含两个子图的图形
    fig = plt.figure(figsize=(12, 5))
    ax1 = fig.add_subplot(121, projection='3d')  # 3D plot
    ax2 = fig.add_subplot(122)  # 2D plot

    # 绘制第一个图形
    for i in range(5):
        for ii in range(collect_capacity):
            rnn_state_trajectory_pca = pca.transform(rnn_state_trajectory_set[i*collect_capacity + ii])
            ax1.plot(rnn_state_trajectory_pca[-seq_len*2:-1, 0], rnn_state_trajectory_pca[-seq_len*2:-1, 1], rnn_state_trajectory_pca[-seq_len*2:-1, 2], c=colors[i], alpha=0.5)

    # 绘制第二个图形
    ax2.plot(tracelet_slices1[0, :, 0], tracelet_slices1[0, :, 1], c=colors[0])
    ax2.plot(tracelet_slices2[0, :, 0], tracelet_slices2[0, :, 1], c=colors[1])
    ax2.plot(tracelet_slices3[0, :, 0], tracelet_slices3[0, :, 1], c=colors[2])
    ax2.plot(tracelet_slices4[0, :, 0], tracelet_slices4[0, :, 1], c=colors[3])
    ax2.plot(tracelet_slices5[0, :, 0], tracelet_slices5[0, :, 1], c=colors[4])

    # 调整子图之间的间距
    plt.subplots_adjust(wspace=0.3)

    # 显示图形
    plt.show()

    # 保存所有数据
    # 1. rnn_state_trajectory_set
    # 2. tracelet_slices1-5
    # 3. seq_len
    # 4. collect_capacity

    np.savez("./logs/e_r9_data.npz", 
             rnn_state_trajectory_set=rnn_state_trajectory_set, 
             tracelet_slices1=tracelet_slices1, 
             tracelet_slices2=tracelet_slices2, 
             tracelet_slices3=tracelet_slices3, 
             tracelet_slices4=tracelet_slices4, 
             tracelet_slices5=tracelet_slices5, 
             seq_len=seq_len, 
             collect_capacity=collect_capacity
             )

    



if __name__ == "__main__":
    main()